import matplotlib.pyplot as plt
import numpy as np
from utils.data_to_color import ImColor_tp, ImColor_tp_12mins
from utils.mssql import sqlDB
from utils.station_inter import interp
import os
from datetime import datetime, timedelta


def compare(radfile):
    radar_data = np.nan_to_num(np.load(radfile)) / 10
    s_date = os.path.basename(rad_file)[:10]  # + '0000'
    s_min = int(os.path.basename(rad_file).split("_")[1].split('-')[0])
    for i in range(radar_data.shape[0]):
        radar_utc_time = datetime.strptime(s_date, '%Y-%m-%d') + timedelta(minutes=s_min) + timedelta(minutes=i * 6)
        station_bj_time = (radar_utc_time + timedelta(hours=8)).strftime('%Y%m%d%H%M')
        # ftime = str(int(ftime0 - ftime0 % 5 + int((ftime0 % 5 + 2.5) / 5) * 5))
        result = test_DB.select(station_bj_time)
        station_data = interp(result)

        rs_data = np.where(station_data > 0, station_data, radar_data[i])

        plt.subplot(131)
        plt.imshow(ImColor_tp_12mins(radar_data[i].copy()[214: 310, 80:202] * 2))  #
        plt.title(f'radar_data {radar_utc_time}')

        plt.subplot(132)
        plt.imshow(ImColor_tp_12mins(station_data.copy()[214: 310, 80:202] * 2))  #
        plt.title(f'station_data {station_bj_time}')

        plt.subplot(133)
        plt.imshow(ImColor_tp_12mins(rs_data.copy()[214: 310, 80:202] * 2))  #
        plt.title(f'rs_data {station_bj_time}')

        plt.gcf().set_size_inches(16, 16)
        plt.tight_layout()
        plt.show()
        result_file = f"results/{station_bj_time}.jpg"
        # plt.savefig(result_file)
        print(result_file)


if __name__ == '__main__':
    test_DB = sqlDB()
    # 用一个rs变量获取数据

    rad_file = 'data/2020-11-23_432-576.npy'
    compare(rad_file)

    # for file in os.listdir("./data"):
    #     rad_file = f"./data/{file}"
    #     compare(rad_file)
    # exit()

    test_DB.close()
